'''
Created on 13.08.2012

@author: myaut
'''

import string

import xml.etree.ElementTree as etree
from lrex import LrexError
from collections import OrderedDict

# Generated script preamble
preamble = '''#!/usr/local/bin/stap -g

/*
Generated automatically by L-REX
*/

probe begin {
    println("Trace started");
}

%{
#include <linux/ktime.h>

s64 __ktime_ns(void) { 
        struct timespec ts;
        ktime_get_ts(&ts);

        return timespec_to_ns(&ts);
}

%}

function ktime_ns:long() %{
        THIS->__retvalue = __ktime_ns();
%}
'''

# Specification parse error
class LrexSpecError(LrexError):
    def __init__(self, message, el):
        Exception.__init__(self, message)
        
        self.el = el
    
    def dump(self):
        etree.dump(self.el)
        
class Param:
    data_types = {
        'int': ('%d', int),
        'uint': ('%u', int),
        'long': ('%ld', int),
        'ulong':('%ul', int),
        'ptr': ('%p', lambda x: int(x, 16)),
        'str': ('%s', lambda s: s)
    }
    
    def __init__(self):
        self.name = None
        
        self.onreturn = False
    
    def parse(self, el):
        '''Parses XML tag el and fills in Param fields
        
        raises LrexSpecError if values are incorrect'''
        try:
            self.name = el.attrib['name']
        except KeyError:
            raise LrexSpecError('Missing param name', el)
        
        try:
            tname = el.attrib['type']
            t = Param.data_types[tname]
        except KeyError:
            raise LrexSpecError('Invalid or missing param data type %s' % tname, el)
        
        self.fmt = t[0]
        self.read = t[1]
        self.cls = el.attrib['class']
        
        if self.cls not in ['key', 'stat', 'info']:
            raise LrexSpecError('Invalid class attribute %s' % self.cls, el)
        
        self.get = el.attrib['get']
        
        onreturn = el.attrib.get('onreturn')
        
        if onreturn is not None:
            if el.attrib.get('onreturn') == 'true':
                self.onreturn = True
            else:
                raise LrexSpecError('''Param onreturn attribute should be "true" or don't exist''', el)
            
    @staticmethod
    def create_pf():
        pf = Param()
        pf.name = 'pf'
        pf.fmt = '%s'
        pf.get = 'probefunc()'
        pf.read = str
            
        return pf

class ProbePoint:
    def parse(self, el):
        pass
    
    def gen_probe_spec(self):
        '''returns (start spec, return spec)'''
        raise TypeError('ProbePoint is abstract!')
    
    def get_params(self):
        return [] 

class SyscallProbePoint(ProbePoint):
    type_name = 'syscall'
    
    def parse(self, el):
        try:
            self.name = el.attrib['name']
        except KeyError:
            raise LrexSpecError('<pp> has no syscall name specification', el)    
        
    def gen_probe_spec(self):
        start = 'syscall.' + self.name
        end = start + '.return'
        
        return start, end

class KernelProbePoint(ProbePoint):
    type_name = 'kernel'
    
    def parse_fspec(self):
        # SystemTap function specification may look like *_write_end@fs/
        # Remove all wildcards and all symbols after @
        self.fname = string.translate(self.fspec, None, '*?')
        self.fname.lstrip('_')
        try:
            at_idx = self.fname.index('@')
            self.fname = self.fname[:at_idx]
        except ValueError:
            pass
    
    def parse(self, el):
        try:
            self.fspec = el.attrib['function']
        except KeyError:
            raise LrexSpecError('<pp> contains nor start/end neither function specification', el)
        
        self.parse_fspec()
    
    def gen_probe_spec(self):
        start = 'kernel.function("%s").call ?' % self.fspec
        end = 'kernel.function("%s").return ?' % self.fspec
        
        return start, end
    
    def get_params(self):
        pf = Param.create_pf()
        
        return [pf]

class ModuleProbePoint(KernelProbePoint):
    type_name = 'module'
    
    def gen_probe_spec(self): 
        start, end = KernelProbePoint.gen_probe_spec(self)
        
        m_start = ', module("*").function("%s").call ?' % self.fspec
        m_end = ', module("*").function("%s").return ?' % self.fspec
        
        return start + m_start, end + m_end
    
class Probe: 
    START = 0
    END = 1
    
    # Probe template for generator
    template = string.Template('''probe ${probe_spec} { ${prologue}
        printf("${probe_alias};%ld;${param_fmt}\\n", 
            ktime_ns(), ${params}); ${epilogue}
    }
    
    ''')
        
    def __init__(self):       
        self.alias = ''
        
        self.prologue = ''
        self.epilogue = ''
        
        self.pps = []
        self.keys = []
        self.params = OrderedDict()
    
    def parse(self, elem):
        param_els = elem.findall('param')
        pp_els = elem.findall('pp')
        
        self.alias = elem.attrib['alias'] 
        
        if not pp_els:
            raise LrexSpecError('No probe-points defined for probe %s' % self.alias)
        if not param_els:
            raise LrexSpecError('No parameters defined for probe %s' % self.alias)
        
        # Parse probe parameters
        for el in param_els:
            param = Param()
            param.parse(el)
            
            self.params[param.name] = param
            
            if param.cls == 'key':
                self.keys.append(param.name)
        
        pp_types = [SyscallProbePoint, KernelProbePoint, ModuleProbePoint]
        for el in pp_els:
            type_name = el.attrib['type']
            
            for PP in pp_types:
                if PP.type_name == type_name:
                    break
            else:
                raise LrexSpecError('Wrong probe-point type %s' % type_name)
            
            pp = PP()
            pp.parse(el)
            
            self.pps.append(pp)
            
            ext_params = pp.get_params()
            for ext_param in ext_params:
                self.params[ext_param.name] = ext_param
        
        #Read prologue/epilogue
        prol = elem.find('prologue')
        if prol is not None:
            self.prologue = prol.text
        
        epil = elem.find('epilogue')
        if epil is not None:
            self.epilogue = epil.text
        
    def gen_probe(self, se):
        '''Generate systemtap script for probe point
        
        se - generate for start or end probe (START or END)'''
        def _gen_params(d, delim, gen):
            return delim.join([gen(v) for v in d.values()])
        
        p = {'prologue': self.prologue, 'epilogue': self.epilogue}
        d = {} # params_dict
        
        start_pps, end_pps = zip(*[pp.gen_probe_spec() for pp in self.pps])
        
        if se == Probe.START:
            s = '+'
            p['probe_spec'] = ','.join(start_pps)
            d = { k: v for k, v in self.params.items() if not v.onreturn }
        elif se == Probe.END:
            s = '-'
            p['probe_spec'] = ','.join(end_pps)
            d = { k: v for k, v in self.params.items() if k in self.keys }
        
        p['probe_alias'] = s + self.alias
        
        p['param_fmt'] = _gen_params(d, ';', lambda k : k.name + ':' + k.fmt)
        p['params'] = _gen_params(d, ', ', lambda k : k.get)
        
        return Probe.template.substitute(p)

class Subsystem:
    def __init__(self):
        # Contains refs (after resolve phase)
        self.probes = OrderedDict()
        
        self.root = None
        self.finish = None
        self.is_root_subsys = False
        
        self.bind_key = ''
    
    def parse(self, elem, spec):
        try:
            self.name = elem.attrib['name']
        except KeyError:
            raise LrexSpecError('<subsystem> without "name" attribute', elem)
        
        for pel in elem.findall('probe'):
            try:          
                probe = Probe()
                probe.parse(pel)
                
                probe.prologue = spec.prologue + probe.prologue
                
                if probe.alias in self.probes:
                    raise LrexSpecError('probe %s already exists' % probe.alias)
                
                self.probes[probe.alias] = probe
                
                # Determine "root probe" where request begins
                if 'root' in pel.attrib:
                    if self.root is not None:
                        raise LrexSpecError('root probe for subsystem %s already defined' % self.name, pel)
                    self.root = probe
                    self.finish = probe
                
                if 'finish' in pel.attrib:
                    # Overrides root event
                    self.finish = probe
            except KeyError:
                raise LrexSpecError('<probe> tag in <subsystem %s> must contain "name" attribute' % self.name, pel)
        
        if self.root is None:
            raise LrexSpecError('root probe for subsystem %s is not defined' % self.name, elem)
        
        bel = elem.find('bind')
        
        if bel is not None:
            try:
                self.bind_name = bel.attrib['name']
                self.bind_key = bel.attrib['key']
            except KeyError:
                raise LrexSpecError('<bind> tag in <subsystem %s> must contain "key" attribute' % self.name, bel)
        else:
            self.is_root_subsys = True
        
        glue = elem.find('glue')
        
        if glue is None:
            raise LrexSpecError('<glue> for <subsystem %s> is not found' % self.name, elem)
        
        try:
            key = glue.attrib['key']
            fmtstr = glue.attrib.get('fmtstr', '%s')
            
            self.glue = key
            self.glue_fmtstr = '%s: ' + fmtstr
        except KeyError:
            raise LrexSpecError('<bind> tag in <subsystem %s> must contain "key" attribute' % self.name, elem)
        
    def is_root_event(self, event):
        return self.root.alias == event.name
    
    def is_finish_event(self, event):
        return self.finish.alias == event.name

class Spec: 
    def __init__(self):   
        self.subsystems = OrderedDict()
        
        self.prologue = ''
        
        self.root_subsys = []

    def parse(self, name):
        self.et = etree.parse(name)
        
        root = self.et.getroot()
        
        try:
            self.name = root.attrib['name']
        except KeyError:
            raise LrexSpecError('Provide name for <trace>', root)
            
        self.author = root.attrib.get('author', 'Anonymous')
        
        prol = root.find('prologue')
        if prol is not None:
            self.prologue = prol.text
        
        subsys_els = root.findall('subsystem')
        
        for sel in subsys_els:
            subsys = Subsystem()
            subsys.parse(sel, self)
            
            if subsys.name in self.subsystems:
                raise LrexSpecError('subsystem %s already exists' % subsys.name)
            
            self.subsystems[subsys.name] = subsys
            
            # Resolve subsystem
            if not subsys.is_root_subsys:
                if subsys.bind_name in self.subsystems:
                    subsys.bind_subsys = self.subsystems[subsys.bind_name]
                else:
                    raise LrexSpecError("Couldn't bind subsystem, because subsys %s doesn't exist" % subsys.bind_subsys_name)
    
    def find(self, event_name, only_root = False):
        '''Find subsys by root subsys/root event or 
        by any subsys/any event'''
        for subsys in self.subsystems.values():
            for probe in subsys.probes.values():
                if probe.alias == event_name:
                    return subsys, probe
        
        raise LrexSpecError("Couldn't find event %s" % event_name)
    
    def generate(self, outf):
        print >> outf, preamble
        
        for subsys in self.subsystems.values():
            for probe in subsys.probes.values():
                print >> outf, probe.gen_probe(Probe.START)
                print >> outf, probe.gen_probe(Probe.END)